package com.lovely602.ai.alibaba.controller.vectorstore;


import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.time.Duration;
import java.util.List;
import java.util.Map;

import static org.springframework.ai.vectorstore.filter.Filter.ExpressionType.*;

/**
 *
 * @author lizhixing
 */
@Slf4j
@RestController
@RequestMapping(value = "/qdrantVectorStore")
public class QdrantVectorStoreController {

    @Autowired
    private QdrantClient qdrantClient;

    @Autowired
    private QdrantVectorStore vectorStore;

    /**
     * 创建qdrant集合
     *
     */
    @GetMapping(value = "/createCollectionName")
    public void createCollectionName(@RequestParam(value = "collectionName", defaultValue = "vector_store") String collectionName) {

        var vectorParams = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024)
                .build();

        Duration duration = Duration.ofSeconds(3);

        qdrantClient.createCollectionAsync(collectionName, vectorParams, duration);
    }


    /**
     * 添加qdrant数据
     *
     */
    @GetMapping(value = "/add")
    public void add(@RequestParam(value = "text") String text) {

        Document document = Document.builder()
                .text(text)
                .metadata(
                        Map.of("auth", "lzx",
                                "country", "CN",
                                "year", 2025,
                                "isActive", "true")
                )
                .build();

        vectorStore.add(List.of(document));

    }

    /**
     * 简单查询qdrant数据
     *
     */
    @GetMapping(value = "/query")
    public List<Document> query(@RequestParam(value = "query") String query) {

        List<Document> documents = vectorStore.similaritySearch(query);

        return documents;
    }

    /**
     * 高级查询qdrant数据
     */
    @GetMapping(value = "/similaritySearch")
    public List<Document> similaritySearch(@RequestParam(value = "query") String query) {

        Filter.Expression expression = new Filter.Expression(AND,
                new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("CN")),
                new Filter.Expression(AND,
                        new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(2025)),
                        new Filter.Expression(EQ, new Filter.Key("isActive"), new Filter.Value("true"))));

        SearchRequest searchRequest = SearchRequest
                .builder()
                .query(query)
                .topK(3) // 最多取几条
                .similarityThreshold(0.6) // 相似度阈值
                .filterExpression(expression)
                .build();

        List<Document> documents = vectorStore.similaritySearch(searchRequest);

        return documents;

    }

    /**
     * 根据ID删除qdrant数据
     *
     */
    @GetMapping(value = "/deleteById")
    public void deleteById(@RequestParam(value = "id") String id) {

        vectorStore.delete(List.of(id));

    }

    /**
     * 删除qdrant数据
     *
     */
    @GetMapping(value = "/delete")
    public void delete(@RequestParam(value = "year") Integer year) {

        Filter.Expression expression = new Filter.Expression(AND,
                new Filter.Expression(EQ, new Filter.Key("country"), new Filter.Value("CN")),
                new Filter.Expression(AND,
                        new Filter.Expression(GTE, new Filter.Key("year"), new Filter.Value(year)),
                        new Filter.Expression(EQ, new Filter.Key("isActive"), new Filter.Value("true"))));
        vectorStore.delete(expression);

    }


}
